package com.liyunc.demo.detect.service;

import com.liyunc.demo.detect.utils.Constants;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.opencv_core;
import org.bytedeco.javacpp.opencv_objdetect;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.OpenCVFrameConverter;

import java.io.File;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;

import static org.bytedeco.javacpp.opencv_imgcodecs.imwrite;
import static org.bytedeco.javacpp.opencv_imgproc.CV_AA;
import static org.bytedeco.javacpp.opencv_imgproc.CV_BGR2GRAY;
import static org.bytedeco.javacpp.opencv_imgproc.cvtColor;
import static org.bytedeco.javacpp.opencv_imgproc.rectangle;
import static org.bytedeco.javacpp.opencv_imgproc.resize;

/**
 * 检测人脸并保存到硬盘的服务.
 * 为训练模型提供资源.
 */
@Slf4j
public class DetectAndSaveService implements DetectService {

    //每一帧原始图片
    private opencv_core.Mat grabbedImage = null;

    // 原始图片对应的灰度图片对象
    private opencv_core.Mat grayImage = null;

    //分类器
    private opencv_objdetect.CascadeClassifier classifier;

    //转换器
    private final OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();

    // 模型文件下载地址
    private final String modelFileUrl;

    // 存放人脸图片的位置
    private final String basePath;

    //记录图片总数
    private final AtomicInteger num = new AtomicInteger();

    //训练的图片尺寸
    opencv_core.Size size = new opencv_core.Size(Constants.RESIZE_WIDTH, Constants.RESIZE_HEIGHT);

    /**
     * 构造方法，指定模型文件的下载地址.
     *
     * @param modelFileUrl 人脸检测模型地址
     * @param basePath     检测出的人脸小图存放地址
     */
    public DetectAndSaveService(String modelFileUrl, String basePath) {
        this.modelFileUrl = modelFileUrl;
        this.basePath = basePath
                + new SimpleDateFormat("yyyyMMddHHmmss").format(new Date())
                + "-";
    }

    /**
     * 视频采样对象的初始化.
     */
    @Override
    public void init() throws Exception {
        // 打开模型文件
        File file = new File(modelFileUrl);

        // 模型文件下载后的完整地址
        String classifierName = file.getAbsolutePath();

        //根据模型文件实例化分类器
        classifier = new opencv_objdetect.CascadeClassifier(classifierName);
    }

    /**
     * 用原始帧做识别.
     *
     * @param frame 原始帧
     * @return 处理后的帧
     */
    @Override
    public Frame convert(Frame frame) {
        // 由帧转为Mat
        grabbedImage = converter.convert(frame);

        // 灰度Mat，用于检测
        if (null == grayImage) {
            grayImage = DetectService.buildGrayImage(grabbedImage);
        }

        String filePath = basePath + num.incrementAndGet() + Constants.IMG_TYPE;

        // 进行人脸识别，根据结果做处理得到预览窗口显示的帧
        return detectAndSave(classifier, converter, frame, grabbedImage, grayImage, filePath, size);
    }

    /**
     * 程序结束前，释放占有的资源.
     */
    @Override
    public void releaseResource() {
        log.info("releaseResource");
        if (null != grabbedImage) {
            grabbedImage.release();
        }
        if (null != grayImage) {
            grayImage.release();
        }
        if (null != classifier) {
            classifier.close();
        }
    }

    /**
     * 对原始帧处理并保存.
     *
     * @param cascadeClassifier 分类器
     * @param cvFrameConverter  转换工具
     * @param frame             原始帧
     * @param grabbedImage      原始图片的Mat对象
     * @param grayImage         原始图片对应的灰度图片的Mat对象
     * @param filePath          图片的基本路径
     * @param size              训练时要求图片的大小
     * @return 处理之后的帧
     */
    private static Frame detectAndSave(opencv_objdetect.CascadeClassifier cascadeClassifier,
                                       OpenCVFrameConverter.ToMat cvFrameConverter,
                                       Frame frame,
                                       opencv_core.Mat grabbedImage,
                                       opencv_core.Mat grayImage,
                                       String filePath,
                                       opencv_core.Size size) {

        // 当前图片转为灰度图片
        cvtColor(grabbedImage, grayImage, CV_BGR2GRAY);

        // 存放检测结果的容器
        opencv_core.RectVector rectVector = new opencv_core.RectVector();

        // 开始检测
        cascadeClassifier.detectMultiScale(grayImage, rectVector);

        // 检测结果总数
        long count = rectVector.size();

        // 如果没有检测结果就提前返回
        // 当前照片检测出的人脸数如果大于1就提前返回不做处理了，这是因为假定运行程序的时候，
        // 摄像头前面只有一个人，所以如果检测出超过一张人脸，就认为当前照片的检测不准确，
        // 就不再处理当前照片了
        if (count != 1) {
            log.info("detect result error, the result count is [{}]", count);
            return frame;
        }
        // 假设现在是一个人对着摄像头，因为此时检测的结果如果大于1，显然是检测有问题

        opencv_core.Mat faceMat;

        // 如果有检测结果，就根据结果的数据构造矩形框，画在原图上
        // 前面的判断确保了此时只有一个人脸
        opencv_core.Rect rect = rectVector.get(0);

        // 从完整的灰度图片中取得一个矩形小图的Mat对象
        faceMat = new opencv_core.Mat(grayImage, rect);

        // 训练时用到的图片尺寸是固定的，因此这里要调整大小
        resize(faceMat, faceMat, size);

        // 图片的保存位置,imwrite方法可以将Mat以图片的形式保存到硬盘
        imwrite(filePath, faceMat);

        // 人脸的位置信息
        int rectX = rect.x();
        int rectY = rect.y();
        int rectW = rect.width();
        int rectH = rect.height();

        // 在人脸上画矩形
        rectangle(grabbedImage,
                new opencv_core.Point(rectX, rectY),
                new opencv_core.Point(rectX + rectW, rectY + rectH),
                opencv_core.Scalar.RED, 1, CV_AA, 0);

        // 释放检测结果资源
        rectVector.close();

        // 将标注过的图片转化为帧，返回
        return cvFrameConverter.convert(grabbedImage);
    }
}
